import cv2
import numpy as np
import os
import datetime


def get_pixel_coord_in_img(rgb_img, m, n, output_format="all", draw_result="False", save_path="check_calibration"):
    """
    计算标定板rgb图像上特征点像素坐标
    Args:
        rgb_img: rgb图像
        m:标定板横向角点数
        n:标定板竖向角点数
        output_format: 特征点选择：
                            all: 所有角点
                            white:所有白色色块中心点
        draw_result: boolean, 是否将角点检测结果画在图中
        save_path: 当draw_result=True时,将结果图像保存于该路径
    Returns:
        output_format为"all":n*m*2的坐标数据
        Output_format为"white" (n-1)*((m-1)/2)*2的坐标数据

    """
    criteria = (cv2.TERM_CRITERIA_MAX_ITER | cv2.TERM_CRITERIA_EPS, 30, 0.001)
    # assert m > n, print(" m must bigger than n")
    gray = cv2.cvtColor(rgb_img, cv2.COLOR_BGR2GRAY)

    ret, corners = cv2.findChessboardCorners(gray, (m, n), None)

    corners = cv2.cornerSubPix(gray, corners, (5, 5), (-1, -1), criteria)  # 在原角点的基础上寻找亚像素角点，即精确到小数
    corners = np.squeeze(corners)

    corners = np.reshape(corners, (n, m, 2))
    cornerss = np.reshape(corners, (n, m, 2)).astype(np.int16)
    if draw_result:
        # check_path(save_path)
        for i in range(n):
            for j in range(m):
                cv2.circle(rgb_img, cornerss[i, j], 5, [0, 0, 255], -1)
        now = datetime.datetime.now()
        timestr = now.strftime("%Y_%m_%d_%H_%M_%S")
        cv2.imwrite(os.path.join(save_path, timestr + "_cal.jpg"), rgb_img)
    return corners

if __name__ == '__main__':
    image = cv2.imread('F:/bubble_dect/Image_20240108154446821.bmp')
    corners = get_pixel_coord_in_img(image, 3, 3, 'all')
    print(corners)